library(tidyverse)
library(redist)
library(alarmdata)
library(sf)

# Functions from fifty-states
source("R/00_setup.R")
source("https://raw.githubusercontent.com/alarm-redist/fifty-states/main/R/summary_stats.R")
source("https://raw.githubusercontent.com/alarm-redist/fifty-states/main/R/validate.R")
source('https://raw.githubusercontent.com/alarm-redist/fifty-states/main/R/partial_sims.R')

# gen old map and new columns ----
map <- alarm_50state_map('AL')
plans <- alarm_50state_plans('AL')
x <- read_rds('data/AL_cd_final_vtd_20.rds') |>
  as_tibble() |>
  select(GEOID = GEOID20, starts_with('cd_pet'))
map <- map |>
  left_join(x, by = 'GEOID')

map_sub_a <- map |>
  `attr<-`('existing_col', 'cd_pet_a') |>
  filter(!cd_pet_a %in% c(1, 2, 7))

# raceblind ----
set.seed(2020)
plans_rb <- redist_smc(
  map, nsims = 5e3, runs = 2L, ncores = 6,
  counties = county, pop_temper = 0.05
  ) |>
  match_numbers("cd_2020") |>
  group_by(chain) |>
  filter(as.integer(draw) < min(as.integer(draw)) + 2500) |>
  ungroup()
plans_rb <- add_summary_stats(plans_rb, map)

write_rds(plans_rb, "data/alabama_raceblind.rds", compress = 'xz')

# frozen 2 vra ----
set.seed(2020)
plans_2_a <- redist_smc(
  map_sub_a, nsims = 5e3, runs = 2L, ncores = 6,
  counties = county, pop_temper = 0.05
)

pl_freeze_a <- prep_particles(map = map |> `attr<-`('existing_col', 'cd_pet_a'),
               map_plan_list = list(a = list(map = map_sub_a, plans = plans_2_a)),
               uid = GEOID, dist_keep = TRUE, nsims = 1e4)
pl_freeze_a <- pl_freeze_a + 2L
pl_freeze_a[which(map$cd_pet_a == 1),] <- 1L
pl_freeze_a[which(map$cd_pet_a == 2),] <- 2L
pl_freeze_a[which(map$cd_pet_a == 7),] <- 7L

pl_freeze_a <- pl_freeze_a |>
  redist_plans(map = map, algorithm = 'smc', diagnostics = attr(plans_2_a, 'diagnostics'))

pl_freeze_a <- pl_freeze_a |>
  mutate(chain = rep(c(1, 2), each = 35e3))  |>
  group_by(chain) |>
  filter(as.integer(draw) < min(as.integer(draw)) + 2500) |>
  ungroup()

pl_freeze_a <- pl_freeze_a |>
  add_reference(map$cd_2020, 'cd_2020') |>
  add_summary_stats(map)

write_rds(pl_freeze_a, "data/alabama_frozen_districts.rds", compress = 'xz')
